#include "BigramCollection.h"

using std::string;

namespace NSpellCheck {

const int* BigramCollection::SENTENCE_START = reinterpret_cast<const int*>(1);
const int* BigramCollection::SENTENCE_END = reinterpret_cast<const int*>(2);

BigramCollection::BigramCollection(const Trie<int>& dictionary):
    dictionary(dictionary)
{}

BigramCollection::~BigramCollection() {}

void BigramCollection::addBigram(const string& firstWord, const string& secondWord, int frequency) {
    const int* firstIterator = dictionary[firstWord];
    if (!firstIterator) {
        throw string() + firstWord + ": unknown word";
    }
    const int* secondIterator = dictionary[secondWord];
    if (!secondIterator) {
        throw string() + secondWord + ": unknown word";
    }
    addBigram(firstIterator, secondIterator, frequency);
}

void BigramCollection::addBigram(const int* it1, const int* it2, int frequency) {
    bigrams[std::make_pair(it1, it2)] += frequency;
}

int BigramCollection::getFrequency(const int* it1, const int* it2) const {
    std::map<std::pair<const int*, const int*>, int>::const_iterator i =
            bigrams.find(std::make_pair(it1, it2));
    return (i == bigrams.end() ? 0 : i->second);
}

int BigramCollection::getFrequency(const std::string& firstWord, const std::string& secondWord) const {
    const int* firstIterator = dictionary[firstWord];
    if (!firstIterator) {
        return 0;
    }
    const int* secondIterator = dictionary[secondWord];
    if (!secondIterator) {
        return 0;
    }
    return getFrequency(firstIterator, secondIterator);
}

}
